"""
NSGA-II多目标遗传算法调度器
实现基于NSGA-II的多目标优化调度
"""

import numpy as np
from typing import List, Dict, Any, Tuple
from deap import base, creator, tools, algorithms
import random
from datetime import datetime, timedelta
from .base_scheduler import BaseScheduler

class NSGA2Scheduler(BaseScheduler):
    """NSGA-II多目标遗传算法调度器"""
    
    def __init__(self, config: Dict[str, Any] = None):
        super().__init__(config)
        self.population_size = self.config.get('population_size', 100)
        self.generations = self.config.get('generations', 50)
        self.crossover_prob = self.config.get('crossover_prob', 0.9)
        self.mutation_prob = self.config.get('mutation_prob', 0.1)
        
        # 设置DEAP框架
        self._setup_deap()
    
    def _setup_deap(self):
        """设置DEAP框架"""
        # 定义适应度函数（最小化）
        creator.create("FitnessMulti", base.Fitness, weights=(-1.0, -1.0, -1.0, -1.0))
        creator.create("Individual", list, fitness=creator.FitnessMulti)
        
        # 创建工具箱
        self.toolbox = base.Toolbox()
        
        # 注册遗传算子
        self.toolbox.register("evaluate", self._evaluate_individual)
        self.toolbox.register("mate", self._crossover)
        self.toolbox.register("mutate", self._mutate)
        self.toolbox.register("select", tools.selNSGA2)
        
    def optimize(self, 
                 resources: List[Dict], 
                 appointments: List[Dict], 
                 objectives: List[str]) -> Dict[str, Any]:
        """
        执行NSGA-II优化
        
        Args:
            resources: 资源列表
            appointments: 预约列表
            objectives: 优化目标列表
            
        Returns:
            调度结果字典
        """
        self.resources = resources
        self.appointments = appointments
        
        # 生成初始种群
        population = self._generate_initial_population()
        
        # 执行NSGA-II算法
        final_population = algorithms.eaMuPlusLambda(
            population,
            self.toolbox,
            mu=self.population_size,
            lambda_=self.population_size,
            cxpb=self.crossover_prob,
            mutpb=self.mutation_prob,
            ngen=self.generations,
            verbose=True
        )[0]
        
        # 选择最优解
        best_solution = self._select_best_solution(final_population)
        
        # 计算目标值
        objectives_values = self.calculate_objectives(best_solution)
        
        return {
            'solution': best_solution,
            'objectives': objectives_values,
            'algorithm': 'NSGA2',
            'generations': self.generations,
            'population_size': self.population_size
        }
    
    def _generate_initial_population(self) -> List:
        """生成初始种群"""
        population = []
        
        for _ in range(self.population_size):
            individual = self._create_random_individual()
            population.append(individual)
            
        return population
    
    def _create_random_individual(self) -> List:
        """创建随机个体"""
        individual = []
        
        for appointment in self.appointments:
            # 为每个预约随机分配资源
            available_resources = [
                r for r in self.resources 
                if r['status'] == 'available' and r['department_id'] == appointment['department_id']
            ]
            
            if available_resources:
                resource = random.choice(available_resources)
                # 随机选择时间（简化实现）
                start_time = appointment['preferred_time'] + timedelta(
                    minutes=random.randint(-60, 60)
                )
                duration = appointment.get('duration', 30)
                end_time = start_time + timedelta(minutes=duration)
                
                individual.append({
                    'appointment_id': appointment['id'],
                    'resource_id': resource['id'],
                    'start_time': start_time,
                    'end_time': end_time
                })
            else:
                # 如果没有可用资源，标记为未分配
                individual.append({
                    'appointment_id': appointment['id'],
                    'resource_id': None,
                    'start_time': None,
                    'end_time': None
                })
                
        return creator.Individual(individual)
    
    def _evaluate_individual(self, individual: List) -> Tuple[float, float, float, float]:
        """评估个体适应度"""
        # 将个体转换为调度方案
        solution = self._individual_to_solution(individual)
        
        # 计算目标值
        objectives = self.calculate_objectives(solution)
        
        # 返回适应度值（负值表示最小化）
        return (
            -objectives['resource_utilization'],  # 最大化资源利用率
            objectives['patient_wait_time'],       # 最小化等待时间
            -objectives['fairness'],              # 最大化公平性
            objectives['overtime_cost']           # 最小化加班成本
        )
    
    def _individual_to_solution(self, individual: List) -> Dict[str, Any]:
        """将个体转换为调度方案"""
        assignments = {}
        
        for assignment in individual:
            if assignment['resource_id'] is not None:
                resource_id = assignment['resource_id']
                if resource_id not in assignments:
                    assignments[resource_id] = []
                assignments[resource_id].append(assignment)
                
        return {'assignments': assignments}
    
    def _crossover(self, ind1: List, ind2: List) -> Tuple[List, List]:
        """交叉操作"""
        # 单点交叉
        crossover_point = random.randint(1, len(ind1) - 1)
        
        # 交换基因
        ind1[crossover_point:], ind2[crossover_point:] = \
            ind2[crossover_point:], ind1[crossover_point:]
            
        return ind1, ind2
    
    def _mutate(self, individual: List) -> List:
        """变异操作"""
        # 随机选择一些预约重新分配
        mutation_count = max(1, len(individual) // 10)  # 变异10%的基因
        
        for _ in range(mutation_count):
            idx = random.randint(0, len(individual) - 1)
            appointment = self.appointments[idx]
            
            # 重新分配资源
            available_resources = [
                r for r in self.resources 
                if r['status'] == 'available' and r['department_id'] == appointment['department_id']
            ]
            
            if available_resources:
                resource = random.choice(available_resources)
                start_time = appointment['preferred_time'] + timedelta(
                    minutes=random.randint(-60, 60)
                )
                duration = appointment.get('duration', 30)
                end_time = start_time + timedelta(minutes=duration)
                
                individual[idx] = {
                    'appointment_id': appointment['id'],
                    'resource_id': resource['id'],
                    'start_time': start_time,
                    'end_time': end_time
                }
                
        return individual
    
    def _select_best_solution(self, population: List) -> Dict[str, Any]:
        """从最终种群中选择最优解"""
        # 选择第一个非支配解
        for individual in population:
            if individual.fitness.valid:
                return self._individual_to_solution(individual)
                
        # 如果没有有效解，返回第一个个体
        return self._individual_to_solution(population[0])
